package com.sicau.tools;

import jakarta.annotation.Resource;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.document.Document;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.tool.annotation.Tool;
import org.springframework.ai.tool.annotation.ToolParam;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.stereotype.Component;

import java.util.List;
import java.util.stream.Collectors;

@Slf4j
@Component
public class DocumentRetrieverTool {

    @Resource
    private EmbeddingModel embeddingModel;
    @Resource
    private VectorStore vectorStore;

    @Tool(description = "根据查询条件，从数据库中检索相关的文档，目前数据库仅有恋爱相关文档")
    public String retrieve(@ToolParam(description = "需要检索相关内容的查询") String query) {
        log.info("开始检索query: {}", query);
        float[] embed = embeddingModel.embed(query);
        SearchRequest searchRequest = SearchRequest.builder()
                .query(query)
                .similarityThreshold(0.5)
                .topK(5)
                .build();
        List<Document> documents = vectorStore.similaritySearch(searchRequest);
        if (documents.isEmpty()) {
            log.info("没有找到与{}相关的内容", query);
            return "没有找到与" + query + "相关的内容";
        } else {
            String result = documents.stream()
                    .map(Document::getText)
                    .collect(Collectors.joining("\n"));
            log.info("检索结果：{}", result);
            return result;
        }
    }
}
